import math

import cv2
import numpy as np
from pprint import pprint


# 根据阈值进行分割标记
def take_mark(img, mark, threshold):
    h, w = np.shape(img)
    for i in range(h):
        for j in range(w):
            if img[i, j] >= threshold:
                mark[i, j] = 1
            else:
                mark[i, j] = 0


def base_on_threshold0(img, filename='img'):
    # 基于迭代的图像阈值分割
    mark = np.zeros_like(img)
    t0 = np.average(img)
    take_mark(img, mark, t0)
    rmark = (mark == False).astype('int')
    t1 = np.sum(img * mark) / np.sum(mark) * 0.5 + np.sum(img * rmark) / np.sum(rmark) * 0.5
    while abs(t0 - t1) <= 0.5:
        t0 = t1
        take_mark(img, mark, t1)
        rmark = (mark == False).astype('int')
        t1 = np.sum(img * mark) / np.sum(mark) * 0.5 + np.sum(img * rmark) / np.sum(rmark) * 0.5
    take_mark(img, mark, t1)
    cv2.imwrite('./res/exp3/res_' + filename + '.png', mark * 255)
    pprint('threshold: ' + str(t0))


def gaussian_filtering(img, sigma, ksize=5):
    # 进行高斯滤波
    if ksize % 2 != 1:
        pprint("核大小不为基数")
        return

    pad = ksize // 2
    gaussian = np.zeros((ksize, ksize))
    for i in range(-pad, pad + 1):
        for j in range(-pad, pad + 1):
            gaussian[i + pad, j + pad] = 1 / (2 * np.pi * sigma ** 2) * np.exp(- (i ** 2 + j ** 2) / (2 * sigma ** 2))
    gaussian = gaussian / np.sum(gaussian)

    h, w = np.shape(img)
    res = np.zeros((h - 2 * pad, w - 2 * pad))
    for i in range(0, h - 2 * pad):
        for j in range(0, w - 2 * pad):
            res[i, j] = np.sum(img[i:i + ksize, j:j + ksize] * gaussian)

    return res


def calculate_gradient_theta(img):
    # 计算梯度和方向角
    Sx = np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]])
    Sy = np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]])
    h, w = np.shape(img)
    gradient = np.zeros((h - 2, w - 2))
    theta = np.zeros((h - 2, w - 2))

    for i in range(h - 2):
        for j in range(w - 2):
            Gx = np.sum(img[i:i + 3, j:j + 3] * Sx)
            Gy = np.sum(img[i:i + 3, j:j + 3] * Sy)
            # pprint('Gx = ' + str(Gx))
            # pprint('Gy = ' + str(Gy))
            gradient[i, j] = (Gx ** 2 + Gy ** 2) ** 0.5
            if Gx == 0:
                if Gy > 0:
                    theta[i, j] = np.pi / 2
                elif Gy < 0:
                    theta[i, j] = - np.pi / 2
                else:
                    theta[i, j] = 0
            else:
                theta[i, j] = math.atan(Gy / Gx)

    return gradient, theta


def non_maximum_suppression(gradient, theta):
    # 进行非极大值抑制，对应方向上的梯度采用线性插值替代
    nms = np.copy(gradient[1:, 1:])
    h, w = np.shape(gradient)
    for i in range(1, h - 1):
        for j in range(1, w - 1):
            cur = gradient[i, j]
            g1, g2 = 0, 0
            if theta[i, j] >= np.pi / 4:
                weight = 0
                if cur == 0:
                    weight = 1 / cur
                g1 = gradient[i - 1, j] * (1 - weight) + gradient[i - 1, j + 1] * weight
                g2 = gradient[i + 1, j] * (1 - weight) + gradient[i + 1, j - 1] * weight
            elif theta[i, j] >= 0:
                weight = cur
                g1 = gradient[i, j + 1] * (1 - weight) + gradient[i - 1, j + 1] * weight
                g2 = gradient[i, j - 1] * (1 - weight) + gradient[i + 1, j - 1] * weight
            elif theta[i, j] >= -np.pi / 4:
                weight = -cur
                g1 = gradient[i, j + 1] * (1 - weight) + gradient[i + 1, j + 1] * weight
                g2 = gradient[i, j - 1] * (1 - weight) + gradient[i - 1, j - 1] * weight
            else:
                weight = - 1 / cur
                g1 = gradient[i - 1, j] * (1 - weight) + gradient[i - 1, j - 1] * weight
                g2 = gradient[i + 1, j] * (1 - weight) + gradient[i + 1, j + 1] * weight

            if cur < g1 or cur < g2:
                nms[i, j] = 0

    return nms


def hysteresis_thresholding(nms, minVal, maxVal):
    h, w = np.shape(nms)
    thresholding = np.zeros_like(nms)
    # 用双阈值进行阈值化
    visited = np.zeros_like(thresholding).astype(bool)
    for i in range(h):
        for j in range(w):
            if nms[i, j] > maxVal:
                thresholding[i, j] = 2
                visited[i, j] = True
            elif nms[i, j] < minVal:
                thresholding[i, j] = 0
                visited[i, j] = True
            else:
                thresholding[i, j] = 1

    directs = [[-1, -1], [-1, 0], [-1, 1], [0, -1], [0, 1], [1, -1], [1, 0], [1, 1]]

    def dfs(row, col, lpos, lflag):
        if visited[row, col]:
            return
        visited[row, col] = True
        lpos.append([row, col])
        if thresholding[row, col] == 2:
            lflag[0] = False
        for direct in directs:
            x = row + direct[0]
            y = col + direct[1]
            if x < 0 or y < 0 or x >= h or y >= w:
                continue
            dfs(x, y, lpos, lflag)

    # 对弱连通部分进行处理，对于与强连通相连的部分进行保留,对独立的部分进行抑制
    for i in range(h):
        for j in range(w):
            if visited[i, j]:
                continue
            elif thresholding[i, j] == 1:
                flag = [False]
                ls = []
                dfs(i, j, ls, flag)
                if flag[0]:
                    for l in ls:
                        thresholding[l[0], l[1]] = 2
                else:
                    for l in ls:
                        thresholding[l[0], l[1]] = 0

    return thresholding / 2 * 255


if __name__ == '__main__':
    image = cv2.imread('./test/test3_1.jpg', cv2.IMREAD_GRAYSCALE)
    image = image.astype('float64')
    base_on_threshold0(image, 'test3_1')
    gauss_img = gaussian_filtering(image, 1.5)
    grad, theta = calculate_gradient_theta(gauss_img)
    nms = non_maximum_suppression(grad, theta)
    for i in range(1, 6):
        for j in range(i, 12):
            img = hysteresis_thresholding(nms, i, j)
            cv2.imwrite('./res/exp3/Test3_1_' + str(i) + '_' + str(j) + '.jpg', img)

    image = cv2.imread('./test/test3_2.jpg', cv2.IMREAD_GRAYSCALE)
    image = image.astype('float64')
    base_on_threshold0(image, 'test3_2')
    gauss_img = gaussian_filtering(image, 1.5)
    grad, theta = calculate_gradient_theta(gauss_img)
    nms = non_maximum_suppression(grad, theta)
    # 尝试不同的双阈值组合，通过结果来找出合适的双阈值
    for i in range(1, 6):
        for j in range(i, 12):
            img = hysteresis_thresholding(nms, i, j)
            cv2.imwrite('./res/exp3/Test3_2_' + str(i) + '_' + str(j) + '.jpg', img)
